#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# 进程查找器 - 用于查找和筛选目标进程

import psutil
import os
import sys
import argparse
import platform

class ProcessFinder:
    def __init__(self):
        self.system = platform.system()
        # 常见的可能包含Cookie的应用关键字
        self.common_app_keywords = [
            'chrome', 'firefox', 'safari', 'edge', 'browser',  # 浏览器
            'electron', 'node', 'nw',  # Electron/Node应用
            'java', 'javaw',  # Java应用
            'python', 'py',  # Python应用
            'qq', 'wechat', 'weixin', 'tim',  # 聊天应用
            'pinduoduo', 'pdd', 'yangkeduo',  # 拼多多相关
            'taobao', 'tmall', 'alipay',  # 阿里系应用
            'jd', 'jingdong',  # 京东
        ]
    
    def list_all_processes(self):
        """列出所有正在运行的进程"""
        processes = []
        for proc in psutil.process_iter(['pid', 'name', 'username', 'exe', 'cmdline']):
            try:
                processes.append(proc.info)
            except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
                pass
        return processes
    
    def find_processes_by_name(self, name):
        """按名称查找进程"""
        processes = []
        for proc in psutil.process_iter(['pid', 'name', 'username', 'exe', 'cmdline']):
            try:
                if name.lower() in proc.info['name'].lower():
                    processes.append(proc.info)
            except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
                pass
        return processes
    
    def find_likely_cookie_processes(self):
        """查找可能包含Cookie的进程"""
        processes = []
        for proc in psutil.process_iter(['pid', 'name', 'username', 'exe', 'cmdline']):
            try:
                proc_info = proc.info
                name_lower = proc_info['name'].lower()
                
                # 检查进程名是否包含常见应用关键字
                if any(keyword in name_lower for keyword in self.common_app_keywords):
                    processes.append(proc_info)
                    continue
                
                # 检查命令行参数
                if proc_info['cmdline']:
                    cmdline = ' '.join(proc_info['cmdline']).lower()
                    if any(keyword in cmdline for keyword in self.common_app_keywords):
                        processes.append(proc_info)
                        continue
                
                # 检查可执行文件路径
                if proc_info['exe'] and any(keyword in proc_info['exe'].lower() for keyword in self.common_app_keywords):
                    processes.append(proc_info)
            except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
                pass
        return processes
    
    def get_process_details(self, pid):
        """获取进程详细信息"""
        try:
            process = psutil.Process(pid)
            info = {
                'pid': process.pid,
                'name': process.name(),
                'exe': process.exe(),
                'cwd': process.cwd(),
                'cmdline': process.cmdline(),
                'username': process.username(),
                'status': process.status(),
                'create_time': process.create_time(),
                'connections': [],
                'open_files': []
            }
            
            # 尝试获取网络连接
            try:
                connections = process.connections()
                for conn in connections:
                    if conn.status == 'ESTABLISHED':
                        info['connections'].append({
                            'local_addr': f"{conn.laddr.ip}:{conn.laddr.port}",
                            'remote_addr': f"{conn.raddr.ip if conn.raddr else 'None'}:{conn.raddr.port if conn.raddr else 'None'}",
                            'status': conn.status
                        })
            except (psutil.AccessDenied, AttributeError):
                pass
            
            # 尝试获取打开的文件
            try:
                open_files = process.open_files()
                for file in open_files:
                    info['open_files'].append(file.path)
            except (psutil.AccessDenied, AttributeError):
                pass
            
            return info
        except (psutil.NoSuchProcess, psutil.AccessDenied):
            return None
    
    def print_process_info(self, process, detailed=False):
        """打印进程信息"""
        if not process:
            print("进程不存在或无法访问")
            return
        
        print(f"PID: {process['pid']}")
        print(f"名称: {process['name']}")
        
        if detailed:
            if 'exe' in process and process['exe']:
                print(f"可执行文件: {process['exe']}")
            
            if 'cmdline' in process and process['cmdline']:
                print(f"命令行: {' '.join(process['cmdline'])}")
            
            if 'username' in process and process['username']:
                print(f"用户: {process['username']}")
            
            if 'connections' in process and process['connections']:
                print("网络连接:")
                for conn in process['connections']:
                    print(f"  {conn['local_addr']} -> {conn['remote_addr']} ({conn['status']})")
            
            if 'open_files' in process and process['open_files']:
                print("打开的文件:")
                cookie_files = []
                other_files = []
                
                for file_path in process['open_files']:
                    if 'cookie' in file_path.lower() or 'session' in file_path.lower():
                        cookie_files.append(file_path)
                    else:
                        other_files.append(file_path)
                
                if cookie_files:
                    print("  可能的Cookie文件:")
                    for file in cookie_files:
                        print(f"    {file}")
                
                if other_files and len(other_files) <= 10:
                    print("  其他文件:")
                    for file in other_files:
                        print(f"    {file}")
                elif other_files:
                    print(f"  其他文件: {len(other_files)} 个")
        
        print("")

def main():
    parser = argparse.ArgumentParser(description='查找和筛选可能包含Cookie的进程')
    parser.add_argument('-l', '--list', action='store_true', help='列出所有进程')
    parser.add_argument('-n', '--name', type=str, help='按名称查找进程')
    parser.add_argument('-c', '--cookie', action='store_true', help='查找可能包含Cookie的进程')
    parser.add_argument('-p', '--pid', type=int, help='获取特定PID的进程详情')
    parser.add_argument('-d', '--detailed', action='store_true', help='显示详细信息')
    
    args = parser.parse_args()
    
    finder = ProcessFinder()
    
    if args.list:
        processes = finder.list_all_processes()
        print(f"发现 {len(processes)} 个进程:")
        for proc in processes:
            finder.print_process_info(proc, args.detailed)
    
    elif args.name:
        processes = finder.find_processes_by_name(args.name)
        print(f"名称包含 '{args.name}' 的进程 ({len(processes)} 个):")
        for proc in processes:
            finder.print_process_info(proc, args.detailed)
    
    elif args.cookie:
        processes = finder.find_likely_cookie_processes()
        print(f"可能包含Cookie的进程 ({len(processes)} 个):")
        for proc in processes:
            finder.print_process_info(proc, args.detailed)
    
    elif args.pid:
        process = finder.get_process_details(args.pid)
        print(f"PID {args.pid} 的进程详情:")
        finder.print_process_info(process, True)
    
    else:
        parser.print_help()

if __name__ == "__main__":
    main() 